import streamlit as st
import easyocr
import numpy as np
from PIL import Image, ImageDraw, ImageFont
import pyttsx3  # 语音引擎库
import os

# 设置页面配置
st.set_page_config(
    page_title="图片文字识别应用",
    layout="wide"
)

# 设置中文字体
st.markdown("""
<style>
    @import url('https://fonts.googleapis.com/css2?family=Noto+Sans+SC:wght@400;700&display=swap');
    html, body, [class*="css"] {
        font-family: 'Noto Sans SC', sans-serif;
    }
</style>
""", unsafe_allow_html=True)


# 加载OCR模型
@st.cache_resource
def load_ocr_reader(language):
    try:
        reader = easyocr.Reader([language], gpu=False)
        return reader
    except Exception as e:
        st.error(f"模型加载失败: {str(e)}")
        return None


# 在图像上绘制识别结果
def draw_ocr_results(image, results):
    draw = ImageDraw.Draw(image)
    font_size = 20

    try:
        # 尝试加载系统字体
        font = ImageFont.truetype("../file/simhei.ttf", font_size)
    except IOError:
        try:
            # 尝试备用字体
            font = ImageFont.truetype("Arial.ttf", font_size)
        except IOError:
            font = None
            st.warning("未找到中文字体，使用默认字体")

    for (bbox, text, prob) in results:
        # 绘制边界框
        x1, y1 = bbox[0]
        x2, y2 = bbox[1]
        x3, y3 = bbox[2]
        x4, y4 = bbox[3]
        draw.line([(x1, y1), (x2, y2), (x3, y3), (x4, y4), (x1, y1)], fill=(0, 255, 0), width=2)

    return image


# 语音朗读功能（优化版本）
def text_to_speech(text, language):
    if not text.strip():
        st.warning("没有可朗读的文本")
        return None

    try:
        # 创建音频文件夹
        if not os.path.exists("audio"):
            os.makedirs("audio")

        # 统一文件格式为WAV
        audio_file = "audio/output.wav"

        # 初始化语音引擎
        engine = pyttsx3.init()
        engine.setProperty('rate', 170)  # 语速

        # 根据语言选择语音
        voices = engine.getProperty('voices')
        voice_id = None

        if language == 'ch_sim':
            for voice in voices:
                if 'chinese' in voice.id.lower() or 'china' in voice.id.lower():
                    voice_id = voice.id
                    break
        else:
            for voice in voices:
                if 'english' in voice.id.lower() or 'en-us' in voice.id.lower():
                    voice_id = voice.id
                    break

        # 生成语音文件
        engine.save_to_file(text, audio_file)
        engine.runAndWait()

        # 读取语音文件
        with open(audio_file, 'rb') as f:
            audio_bytes = f.read()

        return audio_bytes

    except Exception as e:
        st.error(f"语音生成失败: {str(e)}")
        return None


# 主应用
def main():
    # 初始化会话状态
    if 'recognized_text' not in st.session_state:
        st.session_state.recognized_text = ""

    if 'audio_bytes' not in st.session_state:
        st.session_state.audio_bytes = None

    st.title("图片文字识别应用")
    st.markdown("使用EasyOCR技术从图片中提取文字")

    # 侧边栏设置
    sidebar = st.sidebar
    sidebar.header("设置")

    language = sidebar.selectbox(
        "识别语言",
        ["ch_sim", "en"],
        index=0,
    )
    text_threshold = sidebar.slider(
        "文本置信度阈值",
        min_value=0.5,
        max_value=1.0,
        value=0.5,
        step=0.1,
    )

    sidebar.markdown("---")
    sidebar.subheader("语音与下载")

    # 语音朗读按钮
    if sidebar.button("语音朗读"):
        if not st.session_state.recognized_text.strip():
            sidebar.warning("请先识别图片获取文本")
        else:
            sidebar.info("正在生成语音...")
            audio_bytes = text_to_speech(st.session_state.recognized_text, language)
            if audio_bytes:
                st.session_state.audio_bytes = audio_bytes
                sidebar.success("语音生成成功")
                sidebar.audio(audio_bytes, format='audio/wav')
            else:
                sidebar.error("语音生成失败")

    # 下载按钮
    if st.session_state.recognized_text:
        sidebar.download_button(
            label="下载文本",
            data=st.session_state.recognized_text,
            file_name="extracted_text.txt",
            mime="text/plain"
        )

    # 加载OCR阅读器
    reader = load_ocr_reader(language)

    # 文件上传
    uploaded_file = st.file_uploader("上传图片", type=["png", "jpg", "jpeg", "bmp", "tiff"])

    if uploaded_file is not None:
        try:
            # 显示上传的图片
            image = Image.open(uploaded_file)
            # 转换为numpy数组（解决PIL图像可能的参数兼容性问题）
            image_np = np.array(image)

            # 创建两列布局
            col1, col2 = st.columns(2)

            with col1:
                st.subheader("原始图片")
                # 移除use_container_width参数，改用width=None自适应
                st.image(image_np, caption="上传的图片", width=None)

            # 处理按钮
            if st.button("开始识别"):
                if reader:
                    status_text = st.empty()
                    status_text.text("正在识别文字...")

                    img_np = np.array(image)

                    # 执行OCR识别
                    results = reader.readtext(img_np, detail=1, text_threshold=text_threshold)

                    full_text = "\n".join([text for (_, text, _) in results])
                    st.session_state.recognized_text = full_text

                    with col2:
                        st.subheader("识别结果")

                        # 显示带检测框的图像
                        result_image = draw_ocr_results(image.copy(), results)
                        result_image_np = np.array(result_image)  # 转换为numpy数组
                        st.image(result_image_np, caption="识别结果", width=None)

                        # 显示识别的文本
                        st.subheader("提取文字")
                        st.text_area("识别文本", full_text, height=300)

                    status_text.text("识别完成")
        except Exception as e:
            st.error(f"处理上传文件时出错: {str(e)}")


if __name__ == "__main__":
    main()